package com.itdct.onflow.system.common.websocket

import com.alibaba.fastjson2.toJSONString
import com.itdct.onflow.core.extend.logError
import com.itdct.onflow.core.extend.logInfo
import org.springframework.stereotype.Component
import org.springframework.web.socket.CloseStatus
import org.springframework.web.socket.TextMessage
import org.springframework.web.socket.WebSocketSession
import org.springframework.web.socket.handler.TextWebSocketHandler
import java.time.LocalDateTime
import java.util.concurrent.ConcurrentHashMap

/**
 * @author Zhouwx
 * @date 2025/2/24 17:07:42
 * @version 1.0
 * @description
 */
@Component
open class SysWebSocketHandler : TextWebSocketHandler() {
    // 使用线程安全的 ConcurrentHashMap 保存会话
    private val sessions = ConcurrentHashMap<String, WebSocketSession>()

    // 连接建立时调用
    override fun afterConnectionEstablished(session: WebSocketSession) {
        val accountToken = getAccountTokenFromSession(session)
        sessions[accountToken] = session
        logInfo("New connection established: $accountToken, Total connections: ${sessions.size}")
    }

    // 处理文本消息
    override fun handleTextMessage(session: WebSocketSession, message: TextMessage) {
        try {
            val payload = message.payload
            logInfo("Received message: $payload")

            // 这里可以添加消息处理逻辑
            val response = mapOf(
                "type" to "ack",
                "content" to "Message received at ${LocalDateTime.now()}",
                "original" to payload
            )
            session.sendMessage(TextMessage(response.toJSONString()))
        } catch (e: Exception) {
            logError("Error handling message", e)
        }
    }

    // 连接关闭时调用
    override fun afterConnectionClosed(session: WebSocketSession, closeStatus: CloseStatus) {
        val accountToken = getAccountTokenFromSession(session)
        sessions.remove(accountToken)
        logInfo("Connection closed: $accountToken, Reason: ${closeStatus.reason}")
    }

    // 获取用户ID（根据业务逻辑实现）
    private fun getAccountTokenFromSession(session: WebSocketSession): String {
        return session.attributes["accountToken"]?.toString() ?: "anonymous"
    }

    // 向指定用户发送消息
    fun sendToUser(accountId: Long, token: String, message: Map<String, Any>) {
        val accountToken = "${accountId}_${token}"
        sessions[accountToken]?.let { session ->
            if (session.isOpen) {
                try {
                    session.sendMessage(TextMessage(message.toJSONString()))
                } catch (e: Exception) {
                    logError("Error sending message to $accountToken", e)
                }
            }
        }
    }

    // 广播消息
    fun broadcast(message: Map<String, Any>) {
        val jsonMessage = message.toJSONString()
        sessions.values.forEach { session ->
            if (session.isOpen) {
                try {
                    session.sendMessage(TextMessage(jsonMessage))
                } catch (e: Exception) {
                    logError("Error broadcasting message", e)
                }
            }
        }
    }
}